// Keywords: multispecies, interaction, life history, timescale, parasite, infection, SIR, S-I-R

species all initialize() {
	initializeSLiMModelType("nonWF");
	
	defineConstant("K_MONKEY", 100);             // monkey carrying capacity
	defineConstant("F_MONKEY", 3.0);             // monkey fecundity
	defineConstant("G_MONKEY", 20);              // relative generation timescale
	defineConstant("P_TRANSMISSION", 0.0001);    // probability of transmission
	defineConstant("SUPPRESSION_μ", 30000);      // optimal size for suppression
	defineConstant("SUPPRESSION_σ", 10000);      // width for suppression
	defineConstant("SUPPRESSION_STRENGTH", 0.4); // strength of suppression
	defineConstant("DEATH", 100000);             // level for certain death
}
species monkey initialize() {
	initializeSpecies(tickModulo=G_MONKEY, avatar="🐵", color="tan3");
	initializeSLiMOptions(keepPedigrees=T);
	initializeSex();
	// pedigree IDs are used to identify host-pathogen matches
	// colors: blue == uninfected, yellow -> red == infected, black == dead
}
species pathogen initialize() {
	initializeSpecies(avatar="🦠", color="chartreuse3");
	// pathogen.tag is a counter of the next subpop ID to use
	// subpopulation.tag is the pedigree ID of the host for the subpop
}

ticks all 2: first() {
	if (p1.individualCount == 0)
		stop(monkey.avatar + " extinct");
	if (pathogen.subpopulations.size() == 0)
		stop(pathogen.avatar + " extinct");
}

species monkey reproduction(p1, "F") {
	// monkeys reproduce sexually, non-monogamously
	litterSize = rpois(1, F_MONKEY);
	
	for (i in seqLen(litterSize))
	{
		mate = subpop.sampleIndividuals(1, sex="M");
		subpop.addCrossed(individual, mate);
	}
}
species pathogen reproduction() {
	// the pathogen reproduces clonally
	subpop.addCloned(individual);
}

ticks all 1 early() {
	monkey.addSubpop("p1", K_MONKEY);
	
	// choose initial hosts carrying the infection
	initial_hosts = p1.sampleIndividuals(5, replace=F);
	pathogen.tag = 3;
	
	for (initial_host in initial_hosts)
	{
		// make a pathogen subpop for the host
		pathogen_subpop = pathogen.addSubpop(pathogen.tag, 1);
		pathogen_subpop.tag = initial_host.pedigreeID;
		pathogen.tag = pathogen.tag + 1;
	}
	
	// log some basic output
	logfile = community.createLogFile("host_pathogen_log.csv", logInterval=1);
	logfile.addTick();
	logfile.addSubpopulationSize(p1);
	logfile.addPopulationSize(pathogen);
	logfile.addCustomColumn("host_count", "pathogen.subpopulations.size();");
	logfile.addMeanSDColumns("monkeyAge", "p1.individuals.age;");
}

ticks monkey early() {
	// monkey population regulation
	p1.fitnessScaling = K_MONKEY / p1.individualCount;
}
ticks all early() {
	// horizontal transmission
	if (p1.individualCount > 1)
	{
		pathogenSubpops = pathogen.subpopulations;
		allPathogens = pathogenSubpops.individuals;
		isTransmitted = (rbinom(allPathogens.size(), 1, P_TRANSMISSION) == 1);
		moving = allPathogens[isTransmitted];
		
		for (ind in moving)
		{
			// figure out which host we are in
			hostID = ind.subpopulation.tag;
			currentHost = monkey.individualsWithPedigreeIDs(hostID);
			
			// choose a different host and get its ID
			newHost = p1.sampleIndividuals(1, exclude=currentHost);
			newHostID = newHost.pedigreeID;
			
			// find/create a subpop for the new host and move to it
			newSubpop = pathogenSubpops[pathogenSubpops.tag == newHostID];
			
			if (newSubpop.size() == 0)
			{
				newSubpop = pathogen.addSubpop(pathogen.tag, 0);
				pathogen.tag = pathogen.tag + 1;
				newSubpop.tag = newHostID;
				
				// need to incorporate the new subpop in case it receives another
				pathogenSubpops = c(pathogenSubpops, newSubpop);
			}
			newSubpop.takeMigrants(ind);
		}
	}
}
ticks all early() {
	// disease outcomes: either suppression or mortality
	suppress_scale = SUPPRESSION_STRENGTH / dnorm(0.0, 0.0, SUPPRESSION_σ);
	
	for (pathogenSubpop in pathogen.subpopulations)
	{
		popsize = asFloat(pathogenSubpop.individualCount);
		P_supp = (dnorm(popsize, SUPPRESSION_μ, SUPPRESSION_σ) * suppress_scale);
		P_death = popsize / DEATH;
		
		if (runif(1) < P_supp)
		{
			// the infection is suppressed; host lives, pathogen dies
			pathogenSubpop.removeSubpopulation();
		}
		else if (runif(1) < P_death)
		{
			// the host has died; kill it and its pathogens
			host = monkey.individualsWithPedigreeIDs(pathogenSubpop.tag);
			monkey.killIndividuals(host);
			pathogenSubpop.removeSubpopulation();
		}
	}
}
ticks all early() {
	// color monkeys by their infection level
	pathogenSubpops = pathogen.subpopulations;
	
	for (host in p1.individuals)
	{
		hostID = host.pedigreeID;
		pathogenSubpop = pathogenSubpops[pathogenSubpops.tag == hostID];
		
		if (pathogenSubpop.size() == 1)
		{
			pathogenCount = pathogenSubpop.individualCount;
			hue = max(0.0, 1.0 - pathogenCount / DEATH) * 0.15;
			host.color = rgb2color(hsv2rgb(c(hue, 1, 1)));
		}
		else
			host.color = "cornflowerblue";
	}
}

species monkey survival(p1) {
	// when a monkey dies, any pathogens in it die
	if (!surviving)
	{
		hostID = individual.pedigreeID;
		pathogenSubpops = pathogen.subpopulations;
		pathogenSubpop = pathogenSubpops[pathogenSubpops.tag == hostID];
		pathogenSubpop.removeSubpopulation();
	}
	return NULL;
}

ticks all 2000 late() { }
